/* Copyright (c) 2024 Huawei Technologies Co., Ltd.
 * This file is a part of the CANN Open Software.
 * Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
 * Please refer to the License for details. You may not use this file except in compliance with the License.
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
 * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
 * See LICENSE in the root of the software repository for the full text of the License.
 * ===================================================================================================================*/

#include <gtest/gtest.h>
#include "ascir_ops.h"
#include "inc/graph/symbolic.h"
#include "expression/const_values.h"

#include "ascend_ir/ascend_ir.h"
#include "ascend_ir/utils/ascend_graph_code_dumper.h"

class AscendGraphCodeDumperUT : public testing::Test {
 protected:
  void SetUp() {
    dlog_setlevel(0, 0, 0);
  }

  void TearDown() {
    dlog_setlevel(0, 3, 0);
  }
};

namespace ge {
namespace {
std::string ReadFileContent(const std::string &filePath) {
  std::ifstream file(filePath);
  if (!file.is_open()) {
    std::cerr << "Failed to open file: " << filePath << std::endl;
    return "";
  }

  std::stringstream buffer;
  buffer << file.rdbuf();
  return buffer.str();
}

void ConstructAscGraph(AscGraph &graph) {
  Expression s0 = graph.CreateSizeVar("s0");
  // 也可以创建常量，跟var成图的方式不太一样
  Expression s1 = graph.CreateSizeVar(2);

  Axis &s0_axis = graph.CreateAxis("S0", s0);
  Axis &s1_axis = graph.CreateAxis("S1", s1);

  ascir_op::Data data("data", graph);
  data.attr.sched.exec_order = 0;
  data.attr.sched.axis = {s0_axis.id, s1_axis.id};
  data.y.dtype = ge::DT_FLOAT16;
  data.y.format = ge::FORMAT_ND;
  *data.y.axis = {s0_axis.id, s1_axis.id};
  *data.y.repeats = {s0, s1};
  *data.y.strides = {s1, sym::kSymbolOne};

  ascir_op::Data data1("data1", graph);
  data1.attr.sched.exec_order = 0;
  data1.attr.sched.axis = {s0_axis.id, s1_axis.id};
  data1.y.dtype = ge::DT_FLOAT16;
  data1.y.format = ge::FORMAT_ND;
  *data1.y.axis = {s0_axis.id, s1_axis.id};
  *data1.y.repeats = {s0, s1};
  *data1.y.strides = {s1, sym::kSymbolOne};

  ascir_op::Data data2("data2", graph);
  data2.attr.sched.exec_order = 0;
  data2.attr.sched.axis = {s0_axis.id, s1_axis.id};
  data2.y.dtype = ge::DT_FLOAT16;
  data2.y.format = ge::FORMAT_ND;
  *data2.y.axis = {s0_axis.id, s1_axis.id};
  *data2.y.repeats = {s0, s1};
  *data2.y.strides = {s1, sym::kSymbolOne};

  ascir_op::Add add("add");
  add.x1 = data.y;
  add.x2 = data.y;
  add.attr.sched.exec_order = 1;
  add.attr.sched.axis = {s0_axis.id, s1_axis.id};
  add.y.dtype = ge::DT_FLOAT16;
  add.y.format = ge::FORMAT_ND;
  *add.y.axis = {s0_axis.id, s1_axis.id};
  *add.y.repeats = {s0, s1};
  *add.y.strides = {s1, sym::kSymbolOne};

  ascir_op::Exp exp("exp");
  exp.x = data.y;
  exp.attr.sched.exec_order = 2;
  exp.attr.sched.axis = {s0_axis.id, s1_axis.id};
  exp.y.dtype = ge::DT_FLOAT16;
  exp.y.format = ge::FORMAT_ND;
  *exp.y.axis = {s0_axis.id, s1_axis.id};
  *exp.y.repeats = {s0, s1};
  *exp.y.strides = {s1, sym::kSymbolOne};

  ascir_op::Concat concat("concat");
  concat.x.input_nums_ = 2;
  concat.x[0U] = add.y;
  concat.x[1U] = exp.y;
  concat.attr.sched.exec_order = 3;
  concat.attr.sched.axis = {s0_axis.id, s1_axis.id};
  concat.y.dtype = ge::DT_FLOAT16;
  concat.y.format = ge::FORMAT_ND;
  *concat.y.axis = {s0_axis.id, s1_axis.id};
  *concat.y.repeats = {s0, s1 * ge::Symbol(2)};
  *concat.y.strides = {s1 * ge::Symbol(2), sym::kSymbolOne};

  ascir_op::FakeOpA fake_opa("fake/opa");
  fake_opa.dx.input_nums_ = 2;
  fake_opa.dx[0U] = add.y;
  fake_opa.dx[1U] = exp.y;
  fake_opa.x2 = exp.y; //可选输入
  fake_opa.x3 = data1.y;
  fake_opa.x4 = data2.y;
  fake_opa.attr.sched.exec_order = 4;
  fake_opa.attr.sched.axis = {s0_axis.id, s1_axis.id};
  fake_opa.y.dtype = ge::DT_FLOAT16;
  fake_opa.y.format = ge::FORMAT_ND;
  *fake_opa.y.axis = {s0_axis.id, s1_axis.id};
  *fake_opa.y.repeats = {s0, s1 * ge::Symbol(2)};
  *fake_opa.y.strides = {s1 * ge::Symbol(2), sym::kSymbolOne};

  ascir_op::Output output("output");
  output.attr.sched.exec_order = 4;
  output.x = fake_opa.y;
}
}
TEST_F(AscendGraphCodeDumperUT, test_python_gen) {
  AscGraph graph("test");
  ConstructAscGraph(graph);
  ascir::PythonCodeDumper dumper;
  EXPECT_EQ(dumper.Dump(graph, "./asc_graph_python.py"), SUCCESS);
  const std::string expected_graph_code = R"(# Python code to construct AscGraph
from pyautofuse import ascir
from pyautofuse import Autofuser, AutofuserOptions

graph = ascir.HintGraph("test")
s0 = graph.create_size("s0")
S0 = graph.create_axis("S0", s0)
size_of_S1 = ascir.SizeExpr(2)
S1 = graph.create_axis("S1", size_of_S1)
Data_0 = ascir.ops.Data("data", graph)
Data_0.attr.sched.exec_order = 0
Data_0.attr.sched.axis = [S0, S1]
Data_0.y.dtype = ascir.dtypes.float16
Data_0.y.axis = [S0, S1]
Data_0.y.size = [s0, 2]
Data_0.y.strides = [2, 1]
Data_1 = ascir.ops.Data("data1", graph)
Data_1.attr.sched.exec_order = 0
Data_1.attr.sched.axis = [S0, S1]
Data_1.y.dtype = ascir.dtypes.float16
Data_1.y.axis = [S0, S1]
Data_1.y.size = [s0, 2]
Data_1.y.strides = [2, 1]
Data_2 = ascir.ops.Data("data2", graph)
Data_2.attr.sched.exec_order = 0
Data_2.attr.sched.axis = [S0, S1]
Data_2.y.dtype = ascir.dtypes.float16
Data_2.y.axis = [S0, S1]
Data_2.y.size = [s0, 2]
Data_2.y.strides = [2, 1]
Add_0 = ascir.ops.Add("add")
Add_0.attr.sched.exec_order = 1
Add_0.attr.sched.axis = [S0, S1]
Add_0.y.dtype = ascir.dtypes.float16
Add_0.y.axis = [S0, S1]
Add_0.y.size = [s0, 2]
Add_0.y.strides = [2, 1]
Add_0.x1 = Data_0.y
Add_0.x2 = Data_0.y
Exp_0 = ascir.ops.Exp("exp")
Exp_0.attr.sched.exec_order = 2
Exp_0.attr.sched.axis = [S0, S1]
Exp_0.y.dtype = ascir.dtypes.float16
Exp_0.y.axis = [S0, S1]
Exp_0.y.size = [s0, 2]
Exp_0.y.strides = [2, 1]
Exp_0.x = Data_0.y
Concat_0 = ascir.ops.Concat("concat")
Concat_0.attr.sched.exec_order = 3
Concat_0.attr.sched.axis = [S0, S1]
Concat_0.y.dtype = ascir.dtypes.float16
Concat_0.y.axis = [S0, S1]
Concat_0.y.size = [s0, 4]
Concat_0.y.strides = [4, 1]
Concat_0.x = [Add_0.y, Exp_0.y]
FakeOpA_0 = ascir.ops.FakeOpA("fake/opa")
FakeOpA_0.attr.sched.exec_order = 4
FakeOpA_0.attr.sched.axis = [S0, S1]
FakeOpA_0.y.dtype = ascir.dtypes.float16
FakeOpA_0.y.axis = [S0, S1]
FakeOpA_0.y.size = [s0, 4]
FakeOpA_0.y.strides = [4, 1]
FakeOpA_0.dx = [Add_0.y, Exp_0.y]
FakeOpA_0.x2 = Exp_0.y
FakeOpA_0.x3 = Data_1.y
FakeOpA_0.x4 = Data_2.y
Output_0 = ascir.ops.Output("output")
Output_0.attr.sched.exec_order = 4
Output_0.y.dtype = ascir.dtypes.float32
Output_0.x = FakeOpA_0.y
fuser = Autofuser(AutofuserOptions())
fused_NpuKernel0Graph = fuser.autofuse(graph)
op_proto, tiling_def, host_impl, device_impl = fuser.codegen(graph, fused_NpuKernel0Graph)
)";
  EXPECT_EQ(expected_graph_code, ReadFileContent("./asc_graph_python.py"));
}
}

